import json
import logging
import time
from typing import Dict, Optional, List

from cozepy import COZE_CN_BASE_URL, Coze, TokenAuth, Stream, WorkflowEvent, WorkflowEventType
from cozepy import Message, ChatEventType  # 添加聊天相关的导入

# 配置日志
logger = logging.getLogger(__name__)

# Coze API配置
COZE_API_TOKEN = 'pat_OA0N8ZYfHBPrpBLmPFKf7z9rKCt70Jv1uLtjpt8n92HHm4oJOIoHObOyMarl86Uq'  # 添加Coze API密钥
COZE_API_BASE = COZE_CN_BASE_URL  # 设置Coze API的URL，默认为Coze中国版
WORKFLOW_ID = '7554683835113701391'  # 添加工作流ID
BOT_ID = '7554953673320103971'  # 添加机器人ID

# 初始化Coze客户端
coze = Coze(
    auth=TokenAuth(token=COZE_API_TOKEN),  # 使用TokenAuth进行身份验证
    base_url=COZE_API_BASE  # 使用Coze中国版API
)


def run_coze_workflow(parameters: Dict[str, str]) -> Optional[str]:
    """
    运行Coze工作流并返回生成的图片URL
    Args:
        parameters (Dict[str, str]): 工作流参数，包括img, prompt等
    Returns:
        Optional[str]: 生成的图片URL，如果失败则返回None

    Raises:
        Exception: 工作流执行过程中的各种异常
    """
    try:
        logger.info(f"开始执行Coze工作流，参数: {parameters}")

        # 启动工作流的流式运行
        stream = coze.workflows.runs.stream(
            workflow_id=WORKFLOW_ID,  # 工作流ID
            parameters=parameters  # 工作流参数
        )

        # 处理工作流事件流并返回图片URL
        image_url = _handle_workflow_stream(stream)

        logger.info(f"Coze工作流执行成功，图片URL: {image_url}")
        return image_url

    except Exception as e:
        logger.error(f"执行Coze工作流时出错: {str(e)}")
        raise Exception(f"执行Coze工作流失败: {str(e)}")


def _handle_workflow_stream(stream: Stream[WorkflowEvent]) -> Optional[str]:
    """
    处理工作流事件流，提取生成的图片URL

    Args:
        stream (Stream[WorkflowEvent]): 工作流事件流

    Returns:
        Optional[str]: 生成的图片URL

    Raises:
        Exception: 处理过程中的各种异常
    """
    try:
        for event in stream:
            # 如果是消息事件
            if event.event == WorkflowEventType.MESSAGE:
                logger.debug(f"收到消息事件: {event.message}")

                # 提取图片链接
                if event.message and event.message.content:
                    try:
                        content = json.loads(event.message.content)
                        if "output" in content:
                            image_url = content["output"]
                            return image_url
                    except json.JSONDecodeError as e:
                        logger.warning(f"JSON解析失败: {e}")
                        continue

            # 如果是错误事件
            elif event.event == WorkflowEventType.ERROR:
                logger.error(f"收到错误事件: {event.error}")
                error_message = "未知错误"
                if event.error:
                    error_message = getattr(event.error, 'message', str(event.error))
                raise Exception(f"工作流出错: {error_message}")

            # 如果是中断事件
            elif event.event == WorkflowEventType.INTERRUPT:
                logger.warning(f"收到中断事件: {event.interrupt}")
                # 这里可以根据需要处理中断，当前版本简单跳过

        # 如果遍历完所有事件仍未获取到图片URL
        raise Exception("工作流执行完成但未返回图片URL")

    except Exception as e:
        logger.error(f"处理工作流响应时出错: {str(e)}")
        raise Exception(f"处理工作流响应失败: {str(e)}")


# AI提示词生成
class CozeChatBot:
    """
    Coze聊天机器人封装类
    """

    def __init__(self, bot_id: str = BOT_ID, user_id: str = "default_user", conversation_history: List[Message] = None):
        """
        初始化聊天机器人

        Args:
            bot_id (str): 机器人ID
            user_id (str): 用户ID
            conversation_history (List[Message]): 对话历史记录
        """
        self.bot_id = bot_id
        self.user_id = user_id
        self.conversation_history: List[Message] = conversation_history or []

    def chat(self, user_input: str) -> str:
        """
        与机器人进行一轮对话

        Args:
            user_input (str): 用户输入的消息

        Returns:
            str: 机器人的回复
        """
        try:
            # 构建当前用户消息
            user_message = Message.build_user_question_text(user_input)

            # 将当前用户消息添加到对话历史
            self.conversation_history.append(user_message)

            # 机器人回复内容
            bot_response = ""

            # 使用流式接口发起聊天请求
            for event in coze.chat.stream(
                    bot_id=self.bot_id,
                    user_id=self.user_id,
                    additional_messages=self.conversation_history,
            ):
                # 当接收到消息内容增量事件时，累积消息内容
                if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:  # 接收到消息内容增量事件
                    if event.message and event.message.content:  # 确保消息内容不为空
                        bot_response += event.message.content  # 累积消息内容

                # 对话完成事件
                if event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
                    logger.info(f"token usage: {event.chat.usage.token_count}")

            # 构建机器人回复消息并添加到历史记录
            bot_message = Message.build_assistant_answer(bot_response)  # 构建机器人回复消息
            self.conversation_history.append(bot_message)

            return bot_response

        except Exception as e:
            logger.error(f"与机器人对话时出错: {str(e)}")
            raise Exception(f"机器人对话失败: {str(e)}")

    def chat_stream(self, user_input: str):
        """
        与机器人进行流式对话

        Args:
            user_input (str): 用户输入的消息

        Yields:
            str: 机器人回复的增量内容
        """
        try:
            # 构建当前用户消息
            user_message = Message.build_user_question_text(user_input)

            # 将当前用户消息添加到对话历史
            self.conversation_history.append(user_message)

            # 机器人回复内容
            bot_response = ""

            # 使用流式接口发起聊天请求
            for event in coze.chat.stream(
                    bot_id=self.bot_id,
                    user_id=self.user_id,
                    additional_messages=self.conversation_history,
            ):
                # 当接收到消息内容增量事件时，累积消息内容并yield增量
                if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
                    if event.message and event.message.content:
                        bot_response += event.message.content
                        yield event.message.content

                # 对话完成事件
                if event.event == ChatEventType.CONVERSATION_CHAT_COMPLETED:
                    logger.info(f"token usage: {event.chat.usage.token_count}")

            # 构建机器人回复消息并添加到历史记录
            bot_message = Message.build_assistant_answer(bot_response)
            self.conversation_history.append(bot_message)

        except Exception as e:
            logger.error(f"流式对话时出错: {str(e)}")
            yield f"错误: {str(e)}"

    def get_conversation_history(self) -> List[Message]:
        """
        获取对话历史

        Returns:
            List[Message]: 对话历史记录
        """
        return self.conversation_history

    def set_conversation_history(self, history: List[Message]):
        """
        设置对话历史

        Args:
            history (List[Message]): 对话历史记录
        """
        self.conversation_history = history

    def clear_history(self):
        """
        清除对话历史
        """
        self.conversation_history.clear()


def build_conversation_history_from_db(messages) -> List[Message]:
    """
    从数据库消息记录构建 Coze 对话历史

    Args:
        messages: 数据库中的消息记录

    Returns:
        List[Message]: Coze 格式的对话历史
    """
    history = []
    for msg in messages:
        if msg.message_type == 'user':
            history.append(Message.build_user_question_text(msg.content))
        elif msg.message_type == 'assistant':
            history.append(Message.build_assistant_answer(msg.content))
    return history


# 为了兼容你原来的代码，也可以提供函数式接口
def chat_with_bot(user_input: str, bot_id: str = BOT_ID, user_id: str = "default_user") -> str:
    """
    与Coze机器人对话的便捷函数

    Args:
        user_input (str): 用户输入
        bot_id (str): 机器人ID
        user_id (str): 用户ID

    Returns:
        str: 机器人回复
    """
    # 为每个用户维护一个机器人实例（在实际应用中可能需要更复杂的管理方式）
    # bot = CozeChatBot(bot_id, user_id)

    # 流式接口
    # bot = CozeChatBot(bot_id, user_id)
    bot = CozeChatBot(bot_id, user_id)
    # story = ""
    for chunk in bot.chat_stream(user_input):
        # print(chunk, end='', flush=True)
        # story += chunk
        yield chunk


if __name__ == '__main__':
    print(chat_with_bot("3 岁儿童的小故事"))
